Unguarded access to _fsdp_param_group

fsdp_state._fsdp_param_group is typed as Optional[FSDPParamGroup] in PyTorch's FSDP2 internals — it is None for any FSDP module that does not directly manage parameters (e.g. a container module whose children are individually sharded). Accessing ._reshard_after_forward on it unconditionally will raise AttributeError: 'NoneType' object has no attribute '_reshard_after_forward' in that case.

While in practice fsdp_pre_all_gather is only called for tensors managed by a param group, this assumption is implicit. A guard makes the failure mode explicit and easier to diagnose:

fsdp_state = _get_module_fsdp_state(module)
param_group = fsdp_state._fsdp_param_group
if param_group is None:
    raise RuntimeError(
        "FSDP state for this module has no parameter group; "
        "cannot determine reshard_after_forward."
    )
reshard_after_forward = param_group._reshard_after_forward

Inconsistent import style for TrainingState

TrainingState is imported at the module level (line 10) in float8_tensor.py and at line 13 in mxfp8_tensor.py, but here it's imported lazily inside fsdp_pre_all_gather. While the inline comment about the private API and PyTorch version is valuable, the inconsistency across the three sibling files may confuse readers.

Consider either:

• Moving the TrainingState import to the module level and placing the version comment there (matching the other two files), or
• Adding the same lazy-import pattern and version comment to float8_tensor.py and mxfp8_tensor.py for symmetry.

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

columnwise_usage not derived from training state in non-resharded path

When reshard_after_forward=False, the same all-gathered weight is reused through both forward and backward passes. The code sets:

rowwise_usage = True
columnwise_usage = self._quantizer.columnwise_usage

This means whether columnwise data gets derived locally (and kept) is entirely controlled by the sharded quantizer's setting, not the actual pass. The comment in the previous code explicitly noted that both forms were needed when not resharding. If self._quantizer.columnwise_usage is False (e.g. on an architecture that doesn't need the transpose), columnwise data won't be created and won't be available for the backward pass GEMM.

This matches the pre-existing float8_tensor.py behavior (same pattern there), so it's presumably already validated by the existing usage assumptions — but it would be worth a brief comment here documenting that self._quantizer.columnwise_usage must be True whenever the backward GEMM needs columnwise access for the non-resharding path.